import numpy as np
import copy

class LR_ZFDetector(object):

    def __init__(self, params):

        self.Nr = params['Nr']
        self.Nt = params['Nt']
        self.batch_size = params['BATCH_SIZE']
        self.C = np.expand_dims(params['CONSTELLATION_COMPLEX'],0) # [1,M]
        self.Rx = np.mean(np.square(np.abs(self.C)))

    def __call__(self,x,H,y,sigma2):
        """
        格基规约最小均方误差检测算法
        Input:
        x:发送符号 array(shape=(batch_size, 2*Nt), dtype=float32)
        y:接收信号 array(shape=(batch_size, 2*Nr), dtype=float32)
        H:信道矩阵 array(shape=(batch_size, 2*Nr, 2*Nt), dtype=float32)
        sigma2:噪声方差 array(shape=(batch_size, 1), dtype=float32)
        Output:
        nodes:指标字典
        """
        # 实复变换
        H_c = H[:,0:self.Nr,0:self.Nt] + 1j*H[:,self.Nr:self.Nr*2,0:self.Nt]
        y_c = y[:,0:self.Nr] + 1j*y[:,self.Nr:self.Nr*2]
        x_c = x[:,0:self.Nt] + 1j*x[:,self.Nt:self.Nt*2]

        # Lattice Reduction
        H_,T = self.seysen_algorithm(H_c) # H_ = H_c @ T

        # Projected channel output
        H_hy = self.batch_matvec_mul(np.transpose(np.conjugate(H_),[0, 2, 1]), y_c)

        # Gramian of transposed channel matrix
        H_hH_ = np.matmul(np.transpose(np.conjugate(H_),[0, 2, 1]), H_)
        
        # Inverse Gramian
        H_hH_inv = np.linalg.inv(H_hH_)

        # ZF Detector
        zhat = self.zslicer(self.batch_matvec_mul(H_hH_inv, H_hy))

        # z -> x
        xhat = self.batch_matvec_mul(T,zhat)

        xhat_complex = self.slicer_c(xhat)
        x_complex  = x_c

        totalnum = self.batch_size * self.Nt
        rightnum = self.equal(xhat_complex, x_complex)
        errornum = totalnum - rightnum

        nodes={
            'x_complex':x_complex,
            'xhat_complex':xhat_complex,
            'totalnum':totalnum,
            'errornum':errornum,
        }

        return nodes

    def quantify(self,H,H_dual):
        """
        格基规约 seysen 算法正交性度量函数
        Input:
        H:矩阵 array(shape=(batch_size, Nr, Nt), dtype=complex128)
        H_dual:H对偶矩阵矩阵 array(shape=(batch_size, Nr, Nt), dtype=complex128)
        Output:
        q：array(shape=(batch_size,), dtype=float32)，q 最小值为 Nt
        """
        q_o = np.sum(np.square(np.abs(H)),1)
        q_d = np.sum(np.square(np.abs(H_dual)),1)
        q = np.sum(q_o * q_d,1)
        return q

    def seysen_algorithm(self,H):
        """
        格基规约 seysen 算法， Htilde = H T
        Input:
        H:信道矩阵 array(shape=(batch_size, Nr, Nt), dtype=complex128)
        Output:
        Htilde:规约信道矩阵 array(shape=(batch_size, Nr, Nt), dtype=complex128)
        T:幺模矩阵 array(shape=(batch_size, Nt, Nt), dtype=complex128)
        """
        # 初始
        Htilde = copy.deepcopy(H)
        HhH = np.matmul(np.transpose(np.conjugate(Htilde),[0,2,1]), Htilde)
        HhHinv = np.linalg.inv(HhH)
        Htilde_dual = np.transpose(np.conjugate(np.matmul(HhHinv, np.transpose(np.conjugate(Htilde),[0,2,1]))),[0,2,1])
        T = np.tile(np.expand_dims(np.eye(self.Nt),0),[self.batch_size,1,1]).astype('complex128')
        mask = T.astype('bool')
        S = HhH
        S_dual = HhHinv

        ind = list(range(self.batch_size))

        # print(self.quantify(Htilde,Htilde_dual))

        lambd1 = S_dual / (2*np.reshape(S_dual[mask],[self.batch_size,1,self.Nt]))
        lambd2 = S / (2*np.reshape(S[mask],[self.batch_size,self.Nt,1]))
        vij = np.transpose(lambd1-lambd2,[0,2,1])
        lambdaij = np.round(vij)

        while np.max(np.abs(lambdaij)) > 0:
                
            a = np.transpose(np.matmul(np.reshape(S[mask],[self.batch_size,self.Nt,1]),
                    np.reshape(S_dual[mask],[self.batch_size,1,self.Nt])), [0,2,1])
            deltaij = 2*a*np.real( 2*np.conjugate(lambdaij)*vij - np.square(np.abs(lambdaij)) )
            
            index = np.argmax(np.reshape(np.real(deltaij),[self.batch_size,-1]),1)
            k = index // self.Nt
            l = index % self.Nt
            lambdakl = np.expand_dims(lambdaij[ind,k,l],-1)

            T[ind,:,k] += lambdakl * T[ind,:,l]
            Htilde[ind,:,k] += lambdakl * Htilde[ind,:,l]
            Htilde_dual[ind,:,l] -= np.conjugate(lambdakl) * Htilde_dual[ind,:,k]

            S = np.matmul(np.transpose(np.conjugate(Htilde),[0,2,1]), Htilde)
            S_dual = np.linalg.inv(S)

            lambd1 = S_dual / (2*np.reshape(S_dual[mask],[self.batch_size,1,self.Nt]))
            lambd2 = S / (2*np.reshape(S[mask],[self.batch_size,self.Nt,1]))
            vij = np.transpose(lambd1-lambd2,[0,2,1])
            lambdaij = np.round(vij)

        # print(self.quantify(Htilde,Htilde_dual))

        return Htilde,T

    def zslicer(self,xhat):
        """
        切割函数，软符号变换到最近的硬符号
        xhat:软符号 array(shape=(batch_size, Nt), dtype=float32)
        """
        gap = np.sqrt(self.Rx / 2)
        xhat1 = xhat / gap
        xhat2 = np.round(xhat1)
        xhat3 = xhat2 * gap
        return xhat3

    def batch_matvec_mul(self, A, b):
        """
        矩阵A与矩阵b相乘，其中A.shape=(batch_size, Nr, Nt)
        b.shape = (batch_size, Nt)
        输出矩阵C，C.shape = (batch_size, Nr)
        """
        C = np.matmul(A, np.expand_dims(b, axis=2))

        return np.squeeze(C, -1)

    def slicer_c(self, x):
        """
        软符号转换为硬符号
        Input:
        x: 检测器检测后的信号 Tensor(shape=(batchsize, Nt), dtype=complex64)
        Output:
        xhat_complex: 符号硬估计，Tensor(shape=(batchsize, Nt)， dtype=tf.int32)
        """
        Nt = np.shape(x)[1]
        x_complex = np.reshape(x,[-1, 1])
        indices = np.argmin(np.abs(x_complex - self.C), axis=1).astype(np.int32)
        xhat_complex = self.C[:,indices].T
        xhat_complex = np.reshape(xhat_complex,[-1,Nt])

        return xhat_complex

    def equal(self,x,y):
        """
        Computes the numbers of elements for which x and y are equal
        """
        return np.sum(np.equal(x, y).astype(np.float32))